import argparse
from generation_util import LangchainChatWrapper


def args_init():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_path",
        type=str,
        default="../data/problems_sampled/sampled_competition.json",
        help="ClassEval data",
    )

    parser.add_argument(
        "--greedy",
        type=int,
        default=0,
        help="Whether to generate model results with greedy strategy",
    )

    parser.add_argument(
        "--output_path",
        type=str,
        default="Iteration_4.json",
        help="output file path",
    )

    parser.add_argument(
        "--generation_strategy",
        type=int,
        default=1,
        help="Holistic = 0, Incremental = 1",
    )
    parser.add_argument(
        "--model",
        type=int,
        default=1,
        help="kimi = 1",
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="doubao-1-5-lite-32k-250115",
        # default="deepseek-v3-250324",
        # default="deepseek-r1-distill-qwen-7b-250120",
        # default="deepseek-coder-6.7b-base",
        help="大模型的名称",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.2,
        help="temperature value in generation config",
    )
    parser.add_argument(
        "--max_length",
        type=int,
        default=2048,
        help="max length of model's generation result",
    )

    parser.add_argument(
        "--openai_key",
        type=str,
        # 便携ai
        # default="sk-dWhVLiIGvNHNAXbkc9aoeCnyfkaG1aUpCXWcW2rfKvmw5syf",
        # kimi
        # default="sk-6A4WS0CvqvKNXcuVfLf1EoonbGXAolhxNzBANoUUL8Rp7zDr",
        # 火山方舟
        default="48be6f11-e25c-423e-b599-d9c7f60dbc96",
        help="need openai key if use GPT-3.5 or GPT-4",
    )

    parser.add_argument(
        "--openai_base",
        type=str,
        # 便携ai
        # default="https://api.bianxie.ai/v1",
        # kimi
        # default="https://api.moonshot.cn/v1",
        # 火山方舟
        default="https://ark.cn-beijing.volces.com/api/v3",
        # default="http://localhost:1234/v1",
        help="need openai base if use GPT-3.5 or GPT-4",
    )

    parser.add_argument(
        "--sample",
        type=int,
        default=5,
        help="The number of code samples that are randomly generated for each task.",
    )

    return parser.parse_args()

def main():
    """
    主函数：运行代码生成流程
    """
    args = args_init()
    
    print("=" * 50)
    print("代码生成系统启动")
    print("=" * 50)
    print(f"数据源: {args.data_path}")
    print(f"输出文件: {args.output_path}")
    print(f"模型: {args.model_name}")
    print(f"生成策略: {'迭代' if args.generation_strategy == 1 else '整体'}")
    print(f"生成次数: {args.sample}")
    print(f"温度: {args.temperature}")
    print(f"最大长度: {args.max_length}")
    print("=" * 50)
    
    try:
        # 初始化聊天包装器
        chat_wrapper = LangchainChatWrapper(args)
        
        # 运行生成流程
        results = chat_wrapper.run_generation()
        
        print("\n=" * 50)
        print("代码生成完成！")
        print(f"共处理 {len(results)} 个问题")
        print(f"结果已保存到: {args.output_path}")
        print("=" * 50)
        
    except Exception as e:
        print(f"\n运行过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()